import math
import os
from PIL import Image, ImageDraw
import dashscope

def _calculate_scaled_size(original_height, original_width, vl_high_resolution_images):
    """计算模型内部对图像的缩放尺寸（复刻文档中缩放逻辑）"""
    factor = 28
    min_pixels = factor * factor * 4
    if vl_high_resolution_images:
        max_pixels = 16384 * factor * factor
    else:
        max_pixels = 1280 * factor * factor
    
    h_bar = round(original_height / factor) * factor
    w_bar = round(original_width / factor) * factor
    current_pixels = h_bar * w_bar
    
    if current_pixels > max_pixels:
        beta = math.sqrt((original_height * original_width) / max_pixels)
        h_bar = math.floor(original_height / beta / factor) * factor
        w_bar = math.floor(original_width / beta / factor) * factor
    elif current_pixels < min_pixels:
        beta = math.sqrt(min_pixels / (original_height * original_width))
        h_bar = math.ceil(original_height * beta / factor) * factor
        w_bar = math.ceil(original_width * beta / factor) * factor
    
    return h_bar, w_bar

def _get_original_coords(original_image_path, scaled_bbox, vl_high_resolution_images=False):
    """将模型返回的缩放后图像的bbox坐标映射到原始图像坐标"""
    with Image.open(original_image_path) as img:
        original_height = img.height
        original_width = img.width
    
    h_bar, w_bar = _calculate_scaled_size(
        original_height, original_width, vl_high_resolution_images
    )
    
    scale_h = h_bar / original_height
    scale_w = w_bar / original_width
    
    x1_scaled, y1_scaled, x2_scaled, y2_scaled = scaled_bbox
    x1_original = int(round(x1_scaled / scale_w))
    y1_original = int(round(y1_scaled / scale_h))
    x2_original = int(round(x2_scaled / scale_w))
    y2_original = int(round(y2_scaled / scale_h))
    
    x1_original = max(0, min(x1_original, original_width))
    y1_original = max(0, min(y1_original, original_height))
    x2_original = max(0, min(x2_original, original_width))
    y2_original = max(0, min(y2_original, original_height))
    
    return (x1_original, y1_original, x2_original, y2_original)

def _extract_bbox_from_response(response_text):
    """从模型响应中提取bbox坐标"""
    import json
    import re
    
    # 尝试直接解析JSON
    try:
        data = json.loads(response_text)
        if 'bbox_2d' in data:
            return tuple(data['bbox_2d'])
        # 查找第一个包含bbox_2d的项
        for key, value in data.items():
            if isinstance(value, dict) and 'bbox_2d' in value:
                return tuple(value['bbox_2d'])
    except:
        pass
    
    # 使用正则表达式提取坐标
    bbox_pattern = r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]'
    match = re.search(bbox_pattern, response_text)
    if match:
        return tuple(map(int, match.groups()))
    
    raise ValueError(f"无法从响应中提取bbox坐标: {response_text}")

def mark_image_with_bbox(image_path, description, api_key="sk-0e687ddcf0164a6fb66c1096447223c4", output_suffix="_marked"):
    """
    在图片上标记指定描述对应的区域
    
    参数:
        image_path: 图片路径（支持本地路径和file://格式）
        description: 描述文本，用于定位要标记的区域
        api_key: 阿里百炼API密钥（默认值已设置）
        output_suffix: 输出文件后缀
    
    返回:
        输出图片的路径
    """
    # 处理文件路径
    if image_path.startswith("file://"):
        local_path = image_path.replace("file://", "")
        file_uri = image_path
    else:
        local_path = image_path
        file_uri = f"file://{image_path}"
    
    # 构建提示词
    prompt = f"找到{description}，以 JSON 格式输出其bbox 的坐标，不要输出```json```代码段。"
    
    # 调用多模态模型
    messages = [
        {
            "role": "system",
            "content": [{"text": "You are a helpful assistant."}]
        },
        {
            "role": "user",
            "content": [
                {"image": file_uri},
                {"text": prompt}
            ]
        }
    ]
    
    response = dashscope.MultiModalConversation.call(
        api_key=api_key,
        model='qwen-vl-max-latest',
        vl_high_resolution_images=False,
        messages=messages
    )
    
    # 提取响应文本
    response_text = response.output.choices[0].message.content[0]["text"]
    print(f"模型响应: {response_text}")
    
    # 提取bbox坐标
    scaled_bbox = _extract_bbox_from_response(response_text)
    print(f"提取的缩放坐标: {scaled_bbox}")
    
    # 映射到原始图像坐标
    original_bbox = _get_original_coords(local_path, scaled_bbox, vl_high_resolution_images=False)
    print(f"原始图像坐标: {original_bbox}")
    
    # 在原图上画蓝色框
    with Image.open(local_path) as img:
        draw = ImageDraw.Draw(img)
        x1, y1, x2, y2 = original_bbox
        draw.rectangle([x1, y1, x2, y2], outline='blue', width=3)
        
        # 生成输出文件名
        base_name = os.path.splitext(local_path)[0]
        ext = os.path.splitext(local_path)[1]
        output_path = f"{base_name}{output_suffix}{ext}"
        
        img.save(output_path)
        print(f"标记完成，保存为: {output_path}")
        
        return output_path

def main():
    """命令行入口"""
    import sys
    
    if len(sys.argv) < 3:
        print("用法: python vl_bbox_marker.py <图片路径> <描述文本>")
        print("示例: python vl_bbox_marker.py C:/Users/serap/Pictures/1.jpg '十一、公司业务主管楼下方的可签字区域'")
        sys.exit(1)
    
    image_path = sys.argv[1]
    description = sys.argv[2]
    
    try:
        output_path = mark_image_with_bbox(image_path, description)
        print(f"\n✅ 成功！输出文件: {output_path}")
    except Exception as e:
        print(f"❌ 错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()